//命令行解析
use odbc_api::sys::{AttrConnectionPooling, AttrCpMatch};
use odbc_api::{buffers::TextRowSet, Connection, Cursor, Environment,IntoParameter, Error, ResultSetMetadata,parameter::{Blob, BlobRead,WithDataType}};
use std::collections::HashMap;
use std::fmt::Display;
use std::{
    ffi::CStr,
    io::{stdout, Write},
    path::PathBuf,
};
use structopt::StructOpt;
use utf8_builder::Utf8Builder;

use rusty_pool::Builder;
use rusty_pool::ThreadPool;
extern crate base64;

use http::header::HeaderValue;
use std::time::Duration;
use std::{thread, time};

extern crate chrono;
use chrono::offset::Local;
use chrono::DateTime;

use std::sync::{
    atomic::{AtomicI64, Ordering},
    Arc,
};
use ureq::{Agent, AgentBuilder};

use local_ip_address::local_ip;

use std::io;
use tracing::*;
use tracing_subscriber::{self, fmt::time::FormatTime};
use jemallocator;

#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
use serde::Deserialize;
#[derive(Debug, StructOpt)]
#[structopt(name = "rust hive to es8", about = "hive 到 es 的数据传输工具.")]
struct Opt {
    /// The pattern to look for
    /// hive host
    #[structopt(name = "hive host ip", long = "--hive-host")]
    hive_host: String,
    #[structopt(name = "hive db", long = "--hive-db")]
    hive_db: String,
    #[structopt(name = "hive host 端口", long = "--hive-port")]
    hive_port: String,
    #[structopt(name = "hive 用户名", long = "--hive-user")]
    hive_user: String,
    #[structopt(name = "hive 用户密码", long = "--hive-password")]
    hive_password: String,
    #[structopt(name = "hive 查询sql", long = "--hive-sql")]
    hive_sql: String,
    //每批次行数
    #[structopt(name = "每批次多少行", long = "--batch")]
    batch: usize,
    #[structopt(name = "线程池大小", long = "--thread-num")]
    thread_num: usize,

    #[structopt(name = "es地址", long = "--es-url")]
    es_url: String,
    #[structopt(name = "es用户名", long = "--es-user")]
    es_user: String,
    #[structopt(name = "es用户密码", long = "--es-password")]
    es_password: String,
    #[structopt(name = "es索引名", long = "--es-index")]
    es_index: String,

    #[structopt(name = "http超时时间（秒）", long = "--http-timeout")]
    http_timeout: u64,

    #[structopt(name = "错误后重试时间间隔（秒）", long = "--retry-interval")]
    retry_interval: u64,

    #[structopt(name = "重试多少次才放弃", long = "--retry-time")]
    retry_time: u64,

    #[structopt(name = "索引别名", long = "--index-alias")]
    index_alias: String,
    #[structopt(name = "指定字段作为id", long = "--id-name")]
    id_name: String,

    #[structopt(name = "mysql的odbc字符串", long = "--mysql-odbc")]
    mysql_odbc: String,

    #[structopt(name = "多少行显示一次日志", long = "--log-lines")]
    log_lines: usize,


    #[structopt(name = "是否结束错误（1是0否）", long = "--error-exit", default_value = "1")]
    error_exit: u64,
}

/*

测试
RUST_LOG=DEBUG cargo run -- --hive-host=xxxxx --hive-port=10000 --hive-user=hive --hive-password=xxxxx --hive-sql="select * from xxxxxx" --batch=10000 --thread-num=10

*/


#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 初始化并设置日志格式(定制和筛选日志)
    tracing_subscriber::fmt()
        .with_max_level(Level::INFO)
        .with_writer(io::stdout) // 写入标准输出
        .with_ansi(false)  // 如果日志是写入文件，应将ansi的颜色输出功能关掉
        .init();  // 初始化并将SubScriber设置为全局SubScriber

    let start_time = time::SystemTime::now(); //获取开始时间

    //命令行解析
    let opt = Opt::from_args();

    println!("hive_host={}", &opt.hive_host);
    println!("hive_port={}", &opt.hive_port);

    //es bulk 地址
    let es_bulk_url = format!("{}/{}/_bulk", &opt.es_url, &opt.es_index);
    //encoding
    let es_encoding = base64::encode(format!("{}:{}", &opt.es_user, &opt.es_password));

    //线程池
    let thread_pool = Builder::new().core_size(5).max_size(opt.thread_num).build();

    //odbc连接字符串
    let odbc =format!("Driver=/opt/mapr/hiveodbc/lib/64/libmaprhiveodbc64.so;HOST={};PORT={};UID={}; PWD={}; AuthMech=3;RowsFetchedPerBlock={};UseUnicodeSqlCharacterTypes=1",&opt.hive_host,&opt.hive_port,&opt.hive_user,&opt.hive_password,opt.batch.clone());
    //odbc连接环境
    let env = Environment::new().unwrap();

    //成功行数
    let success_count = Arc::new(AtomicI64::new(0));
    //失败行数
    let fail_count = Arc::new(AtomicI64::new(0));

    let mut query_success = true ;
    //odbc连接
    let connection = env.connect_with_connection_string(&odbc)?;


    
    let hive_sql_clone =opt.hive_sql.clone();

    let hive_sqls:Vec<&str> =hive_sql_clone.split(";").collect();

    let hive_sqls_count =hive_sqls.len();

    let mut hive_exec_sql =hive_sql_clone.clone();
    if hive_sqls_count >1 {
    for i in 0..(hive_sqls_count-2)
    {  
       let pre_sql = hive_sqls[i].to_string();

       if pre_sql.starts_with("set ") {
        connection.execute(&pre_sql,()).unwrap_or_default();
       }
    }   
    hive_exec_sql =hive_sqls[hive_sqls_count-1].to_string();
}
     

  
  



    let hive_db_clone =opt.hive_db.clone();
    let mysql_odbc_clone =opt.mysql_odbc.clone();
    match connection.execute(&hive_exec_sql, ())? {
        Some(mut cursor) => {
            // Write the column names to stdout
            let mut headline: Vec<String> = cursor.column_names()?.collect::<Result<_, _>>()?;
            let head_count =headline.len() ;

            for i in 0..head_count{
                let v: Vec<&str> = headline[i].split(".").collect();
                headline[i] = v[v.len() - 1].to_owned();
            }

            log::info!("{:?}", &headline);

            // Use schema in cursor to initialize a text buffer large enough to hold the largest
            // possible strings for each column up to an upper limit of 4KiB.
            let mut buffers = TextRowSet::for_cursor(opt.batch.clone(), &mut cursor, Some(40960))?;
            // Bind the buffer to the cursor. It is now being filled with every call to fetch.
            let mut row_set_cursor = cursor.bind_buffer(&mut buffers)?;

            let mut rowindex :usize =0;

        
            // Iterate over batches
            while let Some(batch) = row_set_cursor.fetch()? {
                //生成_bulk提交字符串
                let mut builder = Utf8Builder::new();


                let start_row_index = rowindex.clone();
                
                

                // Within a batch, iterate over ec_cc_seatvery row
                let rows = batch.num_rows();

                rowindex =rowindex +rows;

                let end_row_index = start_row_index+rows;

                //多少行显示一次日志
                let log_lines_clone = opt.log_lines.clone();
                //指定id字段
                let id_name_clone = opt.id_name.clone();

                let  batch_clone=opt.batch.clone();

                if id_name_clone == "" {
                    for row_index in 0..rows {
                        builder.push_str("{ \"index\" : { }}\n").unwrap_or_default();
                        let mut m: serde_json::Value = serde_json::from_str("{}").unwrap();

                        for col_index in 0..batch.num_cols() {
                            let key =headline[col_index].clone();
                            let value ={
                                match batch.at_as_str(col_index, row_index) {
                                    Ok(s) => s.unwrap_or("").trim().to_owned(),
                                    Err(_) => "".to_owned(),
                                }
                            };
                            if !value.is_empty() {
                               
                                if   (value.starts_with("[") && value.ends_with("]"))   || (value.starts_with("{") && value.ends_with("}")) { 
                                match serde_json::from_str::<serde_json::Value>(&value){
                                    Ok( v)=>{
                                        if(v.is_object()){
                                            m.as_object_mut().unwrap().insert(key.clone(),v);
                                        } else if(v.is_array()) {
                                           
                                            m.as_object_mut().unwrap().insert(key.clone(),v);
                                        } else {
                                            
                                            m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value) );
                                        }
                                    },
                                    Err(_) => {
                                        m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value) );
                                    }
                                }
                            } else {
                                
                                m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value) );
                            }
                            }
                        }
                        builder
                            .push_str(&serde_json::to_string(&m).unwrap_or_default())
                            .unwrap_or_default();
                        builder.push_str("\n").unwrap_or_default();
                        std::mem::drop(m);
                    }
                } else {
                    for row_index in 0..rows {
                        
                        let mut m: serde_json::Value = serde_json::from_str("{}").unwrap();

                        let mut id :String= String::from("");

                        for col_index in 0..batch.num_cols() {

                            let key =headline[col_index].clone();
                            let value ={
                                match batch.at_as_str(col_index, row_index) {
                                    Ok(s) => s.unwrap_or("").to_owned(),
                                    Err(_) => "".to_owned(),
                                }
                            };
                            if !value.is_empty() {
                                if   (value.starts_with("[") && value.ends_with("]"))   || (value.starts_with("{") && value.ends_with("}")) { 
                                    match serde_json::from_str::<serde_json::Value>(&value){
                                        Ok( v)=>{
                                            if(v.is_object()){
                                                m.as_object_mut().unwrap().insert(key.clone(),v);
                                            } else if(v.is_array()) {
                                               
                                                m.as_object_mut().unwrap().insert(key.clone(),v);
                                            } else {
                                                
                                                m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value.clone()) );
                                            }
                                        },
                                        Err(_) => {
                                            m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value.clone()) );
                                        }
                                    }
                                } else {
                                    
                                    m.as_object_mut().unwrap().insert(key.clone(),serde_json::Value::String(value.clone()) );
                                }
                            }
                            if key == id_name_clone {
                                id = value.clone() ;
                            }
                        }

                        builder.push_str("{ \"index\" : { \"_id\":\"").unwrap_or_default();
                        builder.push_str(&id).unwrap_or_default();
                        builder.push_str("\" }}\n").unwrap_or_default();
                        builder
                            .push_str(&serde_json::to_string(&m).unwrap_or_default())
                            .unwrap_or_default();
                        builder.push_str("\n").unwrap_or_default();
                        std::mem::drop(id);
                        std::mem::drop(m);
                    }
                }
                //_bulk 提交字符串
                let result = builder.finalize().unwrap_or_default();
                std::mem::drop(batch);
                let es_bulk_url_clone = es_bulk_url.clone();
                let es_encoding_clone = es_encoding.clone();

                let success_count_clone = success_count.clone();
                let fail_count_clone = fail_count.clone();

                let hive_sql_clone2 =hive_sql_clone.clone();
                let hive_db_clone2= hive_db_clone.clone();
                let mysql_odbc_clone2=mysql_odbc_clone.clone();
                //线程池运行
                thread_pool.execute(move || {
                    loop {
                        let mut c = opt.retry_time.clone(); //重试多少次才放弃
                        match send_to_es(
                            result.clone(),
                            es_bulk_url_clone.clone(),
                            es_encoding_clone.clone(),
                            opt.http_timeout.clone(),
                        ) {
                            Ok(()) => {

                                std::mem::drop(result);
                                //累计成功行数
                                success_count_clone
                                    .fetch_add(rows.clone() as i64, Ordering::Relaxed);

                                let success_lines = success_count_clone.load(Ordering::Relaxed);
                                let fail_lines = fail_count_clone.load(Ordering::Relaxed);

                                
                                if rowindex % log_lines_clone <batch_clone {
                                    log::info!("传输成功{}条,失败{}条", success_lines, fail_lines);
                                }
                                
                                break;
                            }
                            Err(e) => {
                                c = c - 1;
                                if c > 0 {
                                    log::error!("[{}~{}]提交数据库失败:{:?},十秒后重试",start_row_index.clone(),end_row_index.clone() ,&e);
                                } else {
                                    log::error!("提交数据库失败{:?}", &e);
                                    std::mem::drop(result);
                                    //累计失败行数
                                    fail_count_clone
                                        .fetch_add(rows.clone() as i64, Ordering::Relaxed);

                                    let success_lines = success_count_clone.load(Ordering::Relaxed);
                                    let fail_lines = fail_count_clone.load(Ordering::Relaxed);

                                    log::info!("传输成功{}条,失败{}条", success_lines, fail_lines);
                                    
                                    if opt.error_exit ==1 {
                                    let end_time = time::SystemTime::now(); //获取结束始时间              
                                    //mysql日志
                                    save_log(mysql_odbc_clone2, start_time, end_time, success_lines.clone(), fail_lines.clone(), hive_sql_clone2, hive_db_clone2,query_success).unwrap_or_default();
                                    std::process::exit(-3);
                                    }
                                    break;
                                }
                            }
                        }
                        //失败后间隔
                        let retry_interval_seconds =
                            time::Duration::from_millis(opt.retry_interval.clone());
                        thread::sleep(retry_interval_seconds);
                    }
                    ()
                });
            }
        }
        None => {
            log::warn!("hive odbc查询无数据返回.");
            query_success =false;
        }
    }

    //等待线程都结束
    thread_pool.shutdown_join();


    //成功数
    let success_lines = success_count.load(Ordering::Relaxed);
    //失败数
    let fail_lines = fail_count.load(Ordering::Relaxed);

    let end_time = time::SystemTime::now(); //获取结束始时间

    let duration = end_time.duration_since(start_time).unwrap();

    log::info!(
        "执行完成:传输成功{}条,失败{}条,耗时{}秒",
        &success_lines,
        &fail_lines,
        duration.as_secs()
    );

    if query_success && opt.index_alias != "" {
        //先删除已有的
        delete_alias( opt.es_url.clone(), es_encoding.clone(), opt.index_alias.clone(), opt.http_timeout.clone()).unwrap_or_default();
        //新增新的
        add_alias( opt.es_url.clone(), es_encoding.clone(),opt.es_index.clone(), opt.index_alias.clone(), opt.http_timeout.clone()).unwrap_or_default();
    }

    //mysql日志
    save_log(opt.mysql_odbc.clone(), start_time, end_time, success_lines.clone(), fail_lines.clone(), opt.hive_sql.clone(), opt.hive_db.clone(),query_success).unwrap_or_default();
    
    
    if !query_success  {
        std::process::exit(-2);
    }
    
    //有失败行数就当失败
    if fail_lines >0 {
        std::process::exit(-1);
    }

    Ok(())
}

//发送到ES
pub fn send_to_es(
    result: String,
    es_bulk_url: String,
    es_encoding: String,
    time_out: u64,
) -> Result<(), &'static str> {
    //println!("{}",&result);

    /*
    let client = reqwest::blocking::Client::builder()
        .timeout(Duration::from_secs(10))
        .build()?;

    let mut headers = http::header::HeaderMap::new();
    headers.insert(
        http::header::AUTHORIZATION,
        format!("Basic {}", &es_encoding).parse()?,
    );
    headers.insert(
        http::header::CONTENT_TYPE,
        "application/json".parse()?,
    );


    //println!("{:?}",&headers);

    let res = client.post(es_bulk_url).body(result).send()?;
    println!("res={:?}",&res);
    if res.status().is_server_error() {
        log::error!("es bulk 提交失败");
    } else {
        log::info!("es bulk 提交成功");
    }

    log::info!("{:?}", res.status());

    */
    let agent: Agent = ureq::AgentBuilder::new()
        .timeout_read(Duration::from_secs(time_out.clone()))
        .timeout_write(Duration::from_secs(time_out.clone()))
        .build();

    match agent
        .post(&es_bulk_url)
        .set("Authorization", &format!("Basic {}", &es_encoding))
        .set("Content-Type", "application/json")
        .send_string(&result)
    {
        Ok(response) => {
            if response.status() == 200 {
                //log::info!("es bulk 提交成功");
            } else {
                log::error!("{:?}",&response);
                return Err("es bulk 提交失败");
            }
        }
        Err(e) => {
            log::error!("{:?}", &e);
            return Err("es bulk 提交失败");
        }
    }

    std::mem::drop(result);

    Ok(())
}



//删除别名
pub fn delete_alias(
    es_bulk_url: String,
    es_encoding: String,
    alias_name :String,
    time_out: u64,
) -> Result<(), &'static str> {
    
    let agent: Agent = ureq::AgentBuilder::new()
        .timeout_read(Duration::from_secs(time_out.clone()))
        .timeout_write(Duration::from_secs(time_out.clone()))
        .build();

    match agent
        .delete(&format!("{}/_all/_alias/{}",es_bulk_url,alias_name))
        .set("Authorization", &format!("Basic {}", &es_encoding))
        .set("Content-Type", "application/json")
        .call()
    {
        Ok(response) => {
            if response.status() == 200 {
                log::info!("es 删除别名成功");
            } else {
                return Err("es 删除别名失败");
            }
        }
        Err(e) => {
            log::error!("{:?}", &e);
            return Err("es 删除别名失败");
        }
    }

    Ok(())
}


//新增别名
pub fn add_alias(
    es_bulk_url: String,
    es_encoding: String,
    index_name :String,
    alias_name :String,
    time_out: u64,
) -> Result<(), &'static str> {
    
    let agent: Agent = ureq::AgentBuilder::new()
        .timeout_read(Duration::from_secs(time_out.clone()))
        .timeout_write(Duration::from_secs(time_out.clone()))
        .build();

    match agent
        .put(&format!("{}/{}/_alias/{}",es_bulk_url,index_name,alias_name))
        .set("Authorization", &format!("Basic {}", &es_encoding))
        .set("Content-Type", "application/json")
        .call()
    {
        Ok(response) => {
            if response.status() == 200 {
                log::info!("es 新增别名成功");
            } else {
                return Err("es 新增别名失败");
            }
        }
        Err(e) => {
            log::error!("{:?}", &e);
            return Err("es 新增别名失败");
        }
    }

    Ok(())
}


//保存日志
pub fn save_log(mysql_odbc :String ,start_time : time::SystemTime ,end_time : time::SystemTime,success_count :i64 ,fail_count :i64,hive_sql :String ,db :String ,query_success :bool ) -> Result<(), &'static str> {
    
    //odbc连接环境
    let env = Environment::new().unwrap();
    //println!("{}",&mysql_odbc.replace("\\", ""));
    let mut task_id =std::env::var("DOLPHINSCHEDULER_TASKID").unwrap_or_default();
    if task_id =="" {
        task_id =String::from("0");
    }
    match env.connect_with_connection_string(&mysql_odbc.replace("\\", "")){

        
        Ok(connection) => {

            let start_time_local: DateTime<Local> = start_time.into();
            let start_time_local_str :String=format!("{}",start_time_local.format("%Y-%m-%d %T"));

            let end_time_local: DateTime<Local> = end_time.into();
            let end_time_local_str :String=format!("{}",end_time_local.format("%Y-%m-%d %T"));

            let mut ret :String=String::from("0");
            if fail_count>0 {//有失败就失败
                ret=String::from("1");
            }
            if !query_success {//查询失败
                ret=String::from("1");
            }
            let mut msg =format!("成功{}条,失败:{}条",&success_count,&fail_count);
            if !query_success {//查询失败
                msg=String::from("查询失败");
            }


            let ip :String =format!("{:?}",local_ip().unwrap());
            let sql =format!("INSERT INTO ctl_op_log ( log_time, op_type, op_name, op_guid, op_peer, op_bgn_ts, op_end_ts, op_rows_cnt, op_err_code, op_err_mesg,exec_sql,pre_sql,post_sql,srcdb,targetenv,task_id) VALUES (now(), '{}','{}','{}','{}', '{}', '{}', '{}', '{}', '{}','{}','{}','{}','{}','{}','{}')",
    
            &"hive to es",
            &"hive to es",
            &"",
            &ip,
            &start_time_local_str,
            &end_time_local_str,
            &success_count,
            &ret,
            &msg,
            
            &hive_sql.replace("'", "\\'"),
            &"",
            &"",
            &db,
            &std::env::var("CDPENV").unwrap_or_default(),
            &task_id,
            
         );
            connection.execute(&sql,()).unwrap_or_default();
        }
        Err(e)=>{
            log::error!("{:?}",e);
        }
    }


    Ok(())
}